import asyncio
from aiohttp import ClientSession
import json

url = "http://192.168.111.128:8000"

async def completions(model,msg):
    headers = {'content-type': 'application/json'}
    payload = {
        'model': model,
        "messages": msg,
    }

    async with ClientSession() as session:
        async with session.post(url + '/v1/chat/completions', headers=headers, json=payload) as response:
            res = await response.text()
            _r = json.loads(res)
            print(_r)
            return _r['choices'][0]['text'],_r['usage']['prompt_tokens']

async def get_model():
    headers = {'content-type': 'application/json'}

    async with ClientSession() as session:
        async with session.get(url+'/v1/models', headers=headers) as response:
            res = await response.text()
            _r = json.loads(res)
            print(_r)
            return _r['data'][0]['id']

async def main():
    task_list = []
    prompts=[{'role': 'user', 'content': '你好'}]
    model = await asyncio.create_task(get_model())
    print(model)
    task_list = []

    cnt_predict=0
    for msg in prompts:
        task = asyncio.create_task(completions(model,msg))
        cnt_predict +=1
        task_list.append(task)
    await asyncio.wait(task_list, timeout=None)
    num_tokens=0
    for task in task_list:
        ans, tokens= task.result()
        num_tokens += int(tokens)
        print(ans)

res = asyncio.run(main())
print(res)